import torch
import torch.nn as nn


class AlexNet(nn.Module):
    def __init__(self, num_classes=1000):  # 默认分类为1000类（ImageNet数据集）
        super(AlexNet, self).__init__()
        # 特征提取部分 (Feature Extractor)
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),  # 第一层卷积
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),  # 第一层池化
            nn.Conv2d(64, 192, kernel_size=5, padding=2),  # 第二层卷积
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),  # 第二层池化
            nn.Conv2d(192, 384, kernel_size=3, padding=1),  # 第三层卷积
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),  # 第四层卷积
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),  # 第五层卷积
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),  # 第三层池化
        )
        # 分类器部分 (Classifier)
        self.classifier = nn.Sequential(
            nn.Dropout(),  # 防止过拟合
            nn.Linear(256 * 6 * 6, 4096),  # 全连接层1
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),  # 全连接层2
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),  # 输出层
        )

    def forward(self, x):
        x = self.features(x)  # 提取特征
        x = x.view(x.size(0), -1)  # 展平张量
        x = self.classifier(x)  # 分类
        return x


if __name__ == "__main__":
    model = AlexNet()
    print(f'Total parameters: {sum(param.numel() for param in model.parameters())}')
    for name, param in model.named_parameters():
        print(f"层名: {name}, 参数数量: {param.numel()}")
    x = torch.randn((1, 3, 224, 224))
    preds = model(x)
    print('输出shape:', preds.shape)

# Total parameters: 61100840
# 层名: features.0.weight, 参数数量: 23232             nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),  # 第一层卷积
# 层名: features.0.bias, 参数数量: 64
# 层名: features.3.weight, 参数数量: 307200            nn.Conv2d(64, 192, kernel_size=5, padding=2),  # 第二层卷积
# 层名: features.3.bias, 参数数量: 192
# 层名: features.6.weight, 参数数量: 663552            nn.Conv2d(192, 384, kernel_size=3, padding=1),  # 第三层卷积
# 层名: features.6.bias, 参数数量: 384
# 层名: features.8.weight, 参数数量: 884736            nn.Conv2d(384, 256, kernel_size=3, padding=1),  # 第四层卷积
# 层名: features.8.bias, 参数数量: 256
# 层名: features.10.weight, 参数数量: 589824            nn.Conv2d(256, 256, kernel_size=3, padding=1),  # 第五层卷积
# 层名: features.10.bias, 参数数量: 256
# 层名: classifier.1.weight, 参数数量: 37748736         nn.Linear(256 * 6 * 6, 4096),  # 全连接层1
# 层名: classifier.1.bias, 参数数量: 4096
# 层名: classifier.4.weight, 参数数量: 16777216         nn.Linear(4096, 4096),  # 全连接层2
# 层名: classifier.4.bias, 参数数量: 4096
# 层名: classifier.6.weight, 参数数量: 4096000          nn.Linear(4096, num_classes),  # 输出层
# 层名: classifier.6.bias, 参数数量: 1000
# 输出shape: torch.Size([1, 1000])